#!/usr/bin/env python

""" vlan.py Virtual Lan for Emulating a Grid
 Basile Clout, August 2007

This script determines the routing tables of the nodes with respect to the grid configuration.
Usage: ./vlan.py [arguments]
-a, --add_conf	Add the rules (routes and other) to the nodes
-d, --del_conf	Delete these rules
-c, --conf_file	Grid configuration file
-v, --verbose	Display the summary
-h, --help	This help
-f, --fake	Just display the routes, dont send it

Information, help, questions, suggestions, whatever!:
basile.clout@unb.ca
"""

from string import *
from datetime import datetime
import os
import sys
import getopt
import subprocess


class cluster:
    """ Define a cluster structure, a component of a grid """
    def __init__(self, name):
        """ Initialize an instance of cluster with the name of the cluster """
        
        self.name = name                # Name of the cluster
        self.nodes = []                 # List of the nodes in this cluster
        self.adj = []                   # List of the directly connected clusters
        self.routers = []               # List of the routers for this cluster

    def __repr__(self):
        """ Representiation of the cluster """

        sentence = "name = " + str(self.name) +\
                   "\ndirectly connected clusters = " + str(self.adj) +\
                   "\ncontaining the nodes " + str(self.nodes) +\
                   "\nAmong them, some are routers" + str(self.routers)
        
        return sentence
    

class node:
    """ Define a node structure, the component of a cluster """

    def __init__(self, name):

        self.name = name                # Unique identifier for the node
        self.address = "0.0.0.0"        # IP address of the node
        self.adj = []                   # List of external adjacent nodes (for a router)
        self.incluster = ""             # Name of the containing cluster
        
        self.add_conf = ""                 # Order for adding routes in the routing table
        self.del_conf = ""                 # Order for deleting routes in the routing table


    def build_conf(self, nodes, clusters):

        forward = Template('echo $switch > /proc/sys/net/ipv4/ip_forward;\n')
        send_redirects = Template('echo $switch > /proc/sys/net/ipv4/conf/$myif/send_redirects;\n')
        accept_redirects = Template('echo $switch > /proc/sys/net/ipv4/conf/$myif/accept_redirects')
        flush = "ip route flush cache;"
        
        addr, delr = self.build_routes(nodes, clusters)
        self.add_conf += addr
        self.del_conf += delr
        
#        self.add_conf += accept_redirects(switch = 0, myif = "all")
#        self.add_conf += accept_redirects(switch = 0, myif = "eth0")
#        self.add_conf += accept_redirects(switch = 0, myif = "eth1")
        
        if self.adj != []:              # If router
            self.add_conf += forward.substitute(switch = 1)
            self.add_conf += send_redirects.substitute(switch = 0, myif = "all")
            self.add_conf += send_redirects.substitute(switch = 0, myif = "eth0")
            self.add_conf += send_redirects.substitute(switch = 0, myif = "eth1")
            self.del_conf += forward.substitute(switch = 0)
            self.del_conf += send_redirects.substitute(switch = 1, myif = "all")
            self.del_conf += send_redirects.substitute(switch = 1, myif = "eth0")
            self.del_conf += send_redirects.substitute(switch = 1, myif = "eth1")

        self.add_conf += flush
        self.del_conf += flush


    def build_routes(self, nodes, clusters):
        """ Build routes to add and delete in the routing table for this node """

        add_routes = ""
        del_routes = ""
        t = Template('ip route $verb $dest via $gw;\n')

        rip_table = rip(clusters, nodes)

        for _name, dest in nodes.iteritems():

            gate = self.find_gw(dest, nodes, clusters, rip_table.table)
            
            if gate != "":              # If there needs a route AND there is a corresponding gate
                
                add_routes += t.substitute(verb = "add", dest = dest.address, gw = nodes[gate].address)
                del_routes += t.substitute(verb = "del", dest = dest.address, gw = nodes[gate].address)

        return add_routes, del_routes


    def find_gw(self, dest, nodes, clusters, rip):
        """ Find gateway for this node to the destination node """

        if self.name == dest.name:      # This node is the current node
            pass
        elif self.incluster == dest.incluster or dest.name in self.adj: # This node is directly connected to the current node.
            pass
        
        elif self.adj == []:         # This node is NOT a router and the node is in another cluster
            return rip[self.incluster][dest.incluster][1]
        
        else:                           # Current node IS a router
            if self.name != rip[self.incluster][dest.incluster][1]: # If it is a router, but NOT connected to dest.incluster, it behaves like a minion
                return rip[self.incluster][dest.incluster][1]
            else:
                next_cluster = rip[self.incluster][dest.incluster][0] # Find the corresponding cluster
                return rip[next_cluster][self.incluster][1]           # Graph not directed
            
        return ""
            

    def send_order(self, log, cmd = 1):
        """ Send orders adding or deleting routes """
        
        t = Template("ssh $address \"$order\" ")
        if cmd==1:
            string = self.del_conf
        elif cmd==0:
            string = self.add_conf
        else:
            print "Error order to send\n"
            
        order = t.substitute(address = self.address, order = string)

        print "Sending config to " + self.address + " ..."
        log.write("Sending config to " + self.address + " ...\n")
        print order
        log.write(order)

        stdin, stdout, stderr = os.popen3(order)
        output, errput = stdout.read(), stderr.read()

        if errput != "":
            log.writelines("\n" + errput)
            print "Error when sending config to %s (%s)\n\n" % (self.address, self.name)
            log.write("Error when sending config to %s (%s)\n\n" % (self.address, self.name))

        else:
            print "... OK!\n\n"
            log.write("\n ... OK!\n\n")


    def __repr__(self):

        sentence = "node " + self.name +\
                   "\n address " + str(self.address) + \
                   "\n in cluster " + str(self.incluster) +\
                   "\n connected nodes from external cluster (if router) = " + str(self.adj) +\
                   "\n Order add conf = " + str(self.add_conf) +\
                   "\n Order del conf = " + str(self.del_conf)

        return sentence
            

def conf_parse(conf):
    """ Parse a configuration file to fill up the clusters and nodes lists """

    index = 0         # Used for determining Unique name for each node
    new_cluster = 1         # 0 or 1: a blank line shows a new cluster
    nodes = {}                          # List of nodes
    clusters = {}                       # List of clusters

    for line in conf:

        line.strip()                    # Strip spaces in the line
        if line[0]=="#":
            pass
        elif line == "\n":              # A line containing a "newline" is a blank line (after stripping)
            new_cluster = 1
        else:                           # The line defines a new cluster or a new node
            
            parts = line.split(":")
            # TODO: check format regex for clusters and nodes

            if new_cluster == 1:        # The line defines a new cluster


                _cluster = cluster(parts[0].strip())
                _cluster.adj.extend([l.strip() for l in parts[1].split(",")])
                clusters[_cluster.name] = _cluster
                new_cluster = 0

            else:                       # This line defines a new node
                options = parts[1].split(",")

                if options[0].strip() == "":
                    u_name = "m" + str(index) + str(_cluster.name)
                    _node = node(u_name)
                else:
                    _node = node(options[0].strip())

                # TODO: Check regex of the address
                _node.address = parts[0].strip()

                if len(options[1:]) > 0: # If this node is a router
                    _cluster.routers.append(_node.name)
                    
                _node.adj.extend([l.strip() for l in options[1:]])
                _node.incluster = _cluster.name

                _cluster.nodes.append(_node.name)
                nodes[_node.name] = _node

                index += 1

    return clusters, nodes


def which_router_to_this_neighbor(clusters, nodes, c1c, c2c):
    """ Determines which router to use in c1 to reach c2 """

    c1 = clusters[c1c]
    c2 = clusters[c2c]

    if len(c1.routers) == 0:
        print c1.name + "has no routers"
        return ""
    elif len(c1.routers) == 1: # If c1 has only one router (most of the cases)
        return c1.routers[0]
    else:                     # If not, we have to choose the good one
        for r in c1.routers:
            for a in nodes[r].adj:
                if nodes[a].incluster == c2c:
                    return r
    return ""                           # If no router found, error


def find_router_to_cluster(rip, k1, k2):
    """ Return which router in k1 to use to reach k2 """

    


class rip:
    """ Implements a virtual RIP (Routing Information Protocol) protocol """

    def __init__(self, clusters, nodes):
        """ Initialize the dictioary containing the RIP info for the grid """
        self.clusters = clusters
        self.nodes = nodes
        
        self.table = self.build_rip()

    def build_rip(self):
        """ Returns a dictionary containing the RIP infos for all the clusters """

        rip = {}
        
        # Initialization of the entries for each cluster
        for k, c in self.clusters.iteritems():
            rip[k] = self.add_rip_entry(c)

        # Filling the RIP dictionary for each entry
        for k in self.clusters.keys(): # For each entry

            count = len(rip[k])

            while count + 1 < len(self.clusters): # Continue until the number of links for this entry corresponds to the number of clusters in the grid minus 1
                
                for k2 in rip[k].keys(): # Use the entries already existing in k to determine the shortest ways

                    if k2 == k:        # Does not have to happen
                        pass
                    else:
                        rip[k] = self.update_rip_entry(rip, k, k2)

                count = len(rip[k])

        return rip
    

    def add_rip_entry(self, c1):
        """ Initialize an entry corresponding to k1 in the rip dictionary """

        entry = {}                      # Each entry corresponding to one cluster is a dictionary itself

        for c in c1.adj:
            router = which_router_to_this_neighbor(self.clusters, self.nodes, c1.name, c)
            if router == "":
                print "Error in determination of the router from %s\n((%s))\n to %s" % (c1.name, c1, c)
            else:
                entry[c] = (c, router, 1)

        return entry
    

    def update_rip_entry(self, rip, k1, k2):
        """ Update the entry 'k1' of the RIP table 'rip' using the value of the entry k2 """

        for k, desc in rip[k2].iteritems():

            if k == k1:                 # This entry of c2 is c1
                pass;
            else:
                if k in rip[k1].keys(): # If the element in the second entry al ready exist, then check if need to update this entry or not
                    if (rip[k1][k2][2] + desc[2]) < rip[k1][k][2]:
                        rip[k1][k] = (k2, rip[k1][k2][1], rip[k1][k2][2] + desc[2])

                else:                   # Else, insert this entry and update the metric
                    rip[k1][k] = (k2, rip[k1][k2][1], rip[k1][k2][2] + desc[2])

        return rip[k1]


    

def print_summary(nodes, conf_file = "", filename = ""):
    """ Print summary of the orders sent to the nodes """


    if filename == "":

        print "# Vlan summary for " + conf_file 
        print "# " + str(datetime.now()) + "\n"

        for k, n in nodes.iteritems():
            print "\n# Add configuration for %s (%s)" % (n.address, k)
            print n.add_conf
            print "# Del configuration for %s (%s)" % (n.address, k)
            print n.del_conf

        
    else:
        f = file(filename, "w")

        f.write("# Vlan summary for %s \n" % conf_file)
        f.write( "# " + str(datetime.now()) + "\n\n")

        f.write("# Configuration file -------------------------------------")
        conf = file(conf_file,"r")
        f.write(conf.read())
    
        f.write("\n\n# Corresponding orders ----------------------------------- \n\n ")
        for k, n in nodes.iteritems():
            f.write("\n# Add configuration for %s (%s)\n" % (n.address, k))
            f.writelines(n.add_conf)
            f.write("\n# Del configuration for %s (%s)\n" % (n.address, k))
            f.writelines(n.del_conf)

        f.close()
        conf.close()

    
def main():
    """ Main function """

    suf_conf = "conf"
    suf_sum = "sum"
    suf_log = "log"
    conf_file = "vlan.conf"
    switch_mode = 1                          # 1 means delete config, 0 means add
    switch_sum = 0                           # 1 means print sum and log, 0 means just .sum, .log files
    virtue = 0                               # If 1, just print the summary, send nothing
    
    # Parse command line options
    try:
        opts, args = getopt.getopt(sys.argv[1:], "fhadvc:", ["help", "add_conf", "del_conf", "conf_file", "silent", "fake"])
    except getopt.error, msg:
        print msg
        print "For help, use --help"
        sys.exit(2)

    # Process options
    for o, a in opts:

        if o in ("-h", "--help"):
            print __doc__
            sys.exit(0)
        elif o in ("-a", "--add_conf"):
            switch_mode = 0
        elif o in ("-d", "--del_conf"):
            switch_mode = 1
        elif o in ("-c", "conf_file"):
            conf_file = a
        elif o in ("-v","--verbose"):
            switch_sum = 1
        elif o in ("-f", "--fake"):
            virtue = 1

    # Run program
    conf = file(conf_file, "r")
    clusters, nodes = conf_parse(conf)

    sum_file = conf_file.split(".")[0] + "." + suf_sum
    log_file = conf_file.split(".")[0] + "." + suf_log
    
    for k, _node in nodes.iteritems():
        _node.build_conf(nodes, clusters)

    sum = file(sum_file, "w")
        
    if switch_sum == 1:
        print_summary(nodes)
    print_summary(nodes, conf_file, sum_file)
    sum.close()

    if virtue == 0:

        log = file(log_file, "w")
        log.write("# Vlan log for %s \n" % conf_file)
        log.write( "# " + str(datetime.now()) + "\n\n")

        for k, _node in nodes.iteritems():
            _node.send_order(log, switch_mode)
                
        log.close()
        
        
if __name__ == "__main__":
    main()
